import torch
from DEAD.Gradiant.AutoGrad import get_vector_field_component as gvc
from DEAD.Gradiant.AutoGrad import auto_grad
from Utility.RandomSelection import select_random_rows, select_random_rows_for_two_tensors


def loss_sdf(model, data, random_point_num) -> torch.Tensor:
    # 数据解包
    # latent_vector_index,lz和n_slots的维度是[batch_size]
    # burning_surface_coordinate,cavity_region_coordinate和solid_region_coordinate 的维度是[batch_size, n_vertices, 3]
    # cavity_region_sdf和solid_region_sdf 的维度是[batch_size, n_vertices,1]
    latent_vector_index, \
        burning_surface_coordinate, \
        cavity_region_coordinate, \
        cavity_region_sdf, \
        solid_region_coordinate, \
        solid_region_sdf = data

    # 随机选取若干行进行训练
    burning_surface_coordinate = select_random_rows(
        burning_surface_coordinate, random_point_num)
    cavity_region_sdf, cavity_region_coordinate = select_random_rows_for_two_tensors(
        cavity_region_sdf, cavity_region_coordinate, random_point_num)
    solid_region_sdf, solid_region_coordinate = select_random_rows_for_two_tensors(
        solid_region_sdf, solid_region_coordinate, random_point_num)

    # 计算各项损失
    u_pred = model.forward_with_latent_vector_index(
        burning_surface_coordinate, latent_vector_index)
    burning_surface_term = 1e3 * torch.mean(u_pred**2)

    u_pred = model.forward_with_latent_vector_index(
        cavity_region_coordinate, latent_vector_index)
    cavity_term = 1e2 * torch.mean((u_pred - cavity_region_sdf)**2)

    u_pred = model.forward_with_latent_vector_index(
        solid_region_coordinate, latent_vector_index)
    solid_term = 1e2 * torch.mean((u_pred - solid_region_sdf)**2)

    # 计算L1正则化损失
    latent_vectors = model.latent_vectors[latent_vector_index]
    regularization_term = 0.01 * torch.mean(torch.abs(latent_vectors))

    # 总损失
    loss = burning_surface_term + cavity_term + solid_term + regularization_term

    return loss


def loss_eikonal(model, data, random_point_num) -> torch.Tensor:
    # 数据解包
    # latent_vector_index,lz和n_slots的维度是[batch_size]
    # burning_surface_coordinate,cavity_region_coordinate和solid_region_coordinate 的维度是[batch_size, n_vertices, 3]
    # cavity_region_sdf和solid_region_sdf 的维度是[batch_size, n_vertices,1]
    latent_vector_index, \
        burning_surface_coordinate, \
        cavity_region_coordinate, \
        cavity_region_sdf, \
        solid_region_coordinate, \
        solid_region_sdf = data

    # 为了降低显存占用，随机选取若干行进行训练
    burning_surface_coordinate = select_random_rows(
        burning_surface_coordinate, random_point_num)
    cavity_region_sdf, cavity_region_coordinate = select_random_rows_for_two_tensors(
        cavity_region_sdf, cavity_region_coordinate, random_point_num)
    solid_region_sdf, solid_region_coordinate = select_random_rows_for_two_tensors(
        solid_region_sdf, solid_region_coordinate, random_point_num)
    #cavity_region_coordinate = select_random_rows(
    #    cavity_region_coordinate, random_point_num)
    #solid_region_coordinate = select_random_rows(
    #    solid_region_coordinate, random_point_num)

    # 设置梯度
    # bottom_theta_boundary.requires_grad = True
    # top_theta_boundary.requires_grad = True
    burning_surface_coordinate.requires_grad = True
    cavity_region_coordinate.requires_grad = True
    solid_region_coordinate.requires_grad = True

    # u程函方程 1-ux^2-uy^2-uz^2=0
    # 燃面程函项
    u_pred = model.forward_with_latent_vector_index(
        burning_surface_coordinate, latent_vector_index)
    du_dxyz = auto_grad(u_pred, burning_surface_coordinate, 0)
    eikonal_surface_term = torch.mean(
        (1.0-(gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2))**2)

    # 燃面边界条件项
    burning_surface_term = 1e3*torch.mean(u_pred**2)

    # cavity区域项
    u_pred = model.forward_with_latent_vector_index(
        cavity_region_coordinate, latent_vector_index)
    cavity_term = 1e2*torch.mean((u_pred-cavity_region_sdf)**2)

    du_dxyz = auto_grad(u_pred, cavity_region_coordinate, 0)
    eikonal_cavity_term = torch.mean(
        (1.0-(gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2))**2)

    # solid区域项
    u_pred = model.forward_with_latent_vector_index(
        solid_region_coordinate, latent_vector_index)
    solid_term = 1e2*torch.mean((u_pred-solid_region_sdf)**2)

    du_dxyz = auto_grad(u_pred, solid_region_coordinate, 0)
    eikonal_solid_term = torch.mean(
        (1.0-(gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2))**2)

    # 计算L1正则化损失
    latent_vectors = model.latent_vectors[latent_vector_index]
    regularization_term = 0.01 * torch.mean(torch.abs(latent_vectors))

    loss = regularization_term+eikonal_surface_term + eikonal_cavity_term + \
        eikonal_solid_term + burning_surface_term +cavity_term + solid_term

    return loss
